/*
 * Copyright (C) 2021 Baidu, Inc. All Rights Reserved.
 */
package com.chukun.rpc.registry.lb;

import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentSkipListMap;

import org.apache.curator.x.discovery.ServiceInstance;

import com.chukun.rpc.core.TinyRpcServiceMetaData;

/**
 * <p><p/>
 *
 * @author chukun
 * @version 1.0
 * @since 1.0-SNAPSHOT
 */
public class ZKConsistentHashLoadBalancer implements ServiceLoadBalancer<ServiceInstance<TinyRpcServiceMetaData>> {

    private final static int VIRTUAL_NODE_SIZE = 10;
    private final static String VIRTUAL_NODE_SPLIT = "#";
    private final static String SERVICE_INSTANCE_KEY_SPLIT = ":";

    @Override
    public ServiceInstance<TinyRpcServiceMetaData> select(List<ServiceInstance<TinyRpcServiceMetaData>> services, int hash) throws Exception {
        ConcurrentSkipListMap<Integer, ServiceInstance<TinyRpcServiceMetaData>> ring = generateConsistentHashRing(services);
        return allocateNode(ring, hash);
    }

    private ServiceInstance<TinyRpcServiceMetaData> allocateNode(ConcurrentSkipListMap<Integer, ServiceInstance<TinyRpcServiceMetaData>> ring, int hash) {
        Map.Entry<Integer, ServiceInstance<TinyRpcServiceMetaData>> entry = ring.ceilingEntry(hash);
        if (entry == null) {
            entry = ring.firstEntry();
        }
        return entry.getValue();
    }

    /**
     * 生成一致性hash算法
     * @param services
     * @return
     */
    private ConcurrentSkipListMap<Integer, ServiceInstance<TinyRpcServiceMetaData>> generateConsistentHashRing(List<ServiceInstance<TinyRpcServiceMetaData>> services) {
        ConcurrentSkipListMap<Integer, ServiceInstance<TinyRpcServiceMetaData>> ring = new ConcurrentSkipListMap<>();
        for (ServiceInstance<TinyRpcServiceMetaData> instance : services) {
            for (int i = 0; i< VIRTUAL_NODE_SIZE; i++) {
                ring.put(String.join(buildServiceInstanceKey(instance), VIRTUAL_NODE_SPLIT, String.valueOf(i)).hashCode(), instance);
            }
        }
        return ring;
    }

    /**
     * 构建注册中心的key
     * @param instance
     * @return
     */
    private String buildServiceInstanceKey(ServiceInstance<TinyRpcServiceMetaData> instance) {
        TinyRpcServiceMetaData metaData = instance.getPayload();
        return String.join(SERVICE_INSTANCE_KEY_SPLIT, metaData.getServiceAddr(), String.valueOf(metaData.getPort()));
    }



}
